/*
 *  Copyright (C) 2015 Australian Institute of Marine Science
 *
 *  Contact: Gael Lafond <g.lafond@aims.org.au>
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package au.gov.aims.concurrent;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.FutureTask;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

public class ExecutorServiceWrapper<T extends FutureTask<V>, V> {
	private ExecutorService pool;

	// Keep track of running tasks, to get the results and call the
	// "onIdle" method when all tasks has been completed.
	private ConcurrentLinkedQueue<T> runningTasks;
	private AtomicInteger nbTotal = null;
	private AtomicInteger nbExecuted = null;

	public ExecutorServiceWrapper() {
		this(Runtime.getRuntime().availableProcessors());
	}

	public ExecutorServiceWrapper(int nbThreads) {
		this.pool = Executors.newFixedThreadPool(nbThreads);
		this.runningTasks = new ConcurrentLinkedQueue<T>();
	}

	// It's recommended to override those 3 methods
	public void onCompleted(List<V> results) {}
	public void onUpdate(T task) {}
	public void onInterruptedException(Exception ex) {
		ex.printStackTrace();
	}
	public void onExecutionException(Exception ex) {
		ex.printStackTrace();
	}

	// Call the execute method from the ExecutorService
	// and keep track of running tasks.
	public void execute(T task) {
		boolean wasIdle = this.runningTasks.isEmpty();
		this.runningTasks.add(task);

		// "Although ExecutorService defines several methods which take a Runnable,
		//  you should generally limit yourself to using the execute method inherited
		//  from Executor.
		// https://docs.oracle.com/javafx/2/api/javafx/concurrent/Task.html
		this.pool.execute(task);

		// Barrier
		// As soon as there is one task to be executed, the wrapper wait for all tasks to be completed,
		// then it call the "onIdle" method.
		// This process can be repeated as often as needed
		// (i.e. new tasks can be added after all the tasks are completed)
		if (wasIdle) {
			this.nbTotal = new AtomicInteger(0);
			this.nbExecuted = new AtomicInteger(0);
			Runnable waitForCompletionTask = new Runnable() {
				public void run() {
					List<V> results = new ArrayList<V>(ExecutorServiceWrapper.this.runningTasks.size());
					T task;
					while ((task = ExecutorServiceWrapper.this.runningTasks.poll()) != null) {
						try {
							// Wait for the task to complete and collect the result.
							// This action freeze the current thread so it has to be ran within an independent thread.
							results.add(task.get());
						} catch (InterruptedException e) {
							ExecutorServiceWrapper.this.onInterruptedException(e);
						} catch (ExecutionException e) {
							ExecutorServiceWrapper.this.onExecutionException(e);
						}
					}
					// Return all results at once when there is no more tasks waiting
					ExecutorServiceWrapper.this.onCompleted(results);
				}
			};
			new Thread(waitForCompletionTask).start();
		}
		this.nbTotal.incrementAndGet();
	}

	public void taskDone(T task) {
		this.nbExecuted.incrementAndGet();
		this.onUpdate(task);
	}

	public final int getDone() {
		return this.nbExecuted.intValue();
	}
	public final int getTotal() {
		return this.nbTotal.intValue();
	}

	// NOTE: Progress may go back if new tasks are added quicker that they are executed.
	public final double getProgress() {
		int done = this.getDone();
		double total = this.getTotal();
		return total > 0 ? done / total : 1.0;
	}

	public void cancel(boolean mayInterrupt) {
		T task;
		while ((task = ExecutorServiceWrapper.this.runningTasks.poll()) != null) {
			task.cancel(mayInterrupt);
		}
	}

	// Shutdown the ExecutorService properly.
	// See: http://docs.oracle.com/javase/7/docs/api/java/util/concurrent/ExecutorService.html
	public void shutdown() {
		this.pool.shutdown(); // Disable new tasks from being submitted
		try {
			// Wait a while for existing tasks to terminate
			if (!this.pool.awaitTermination(60, TimeUnit.SECONDS)) {
				this.pool.shutdownNow(); // Cancel currently executing tasks
				// Wait a while for tasks to respond to being cancelled
				if (!this.pool.awaitTermination(60, TimeUnit.SECONDS)) {
					throw new IllegalStateException("Pool did not terminate");
				}
			}
		} catch (InterruptedException ie) {
			// (Re-)Cancel if current thread also interrupted
			this.pool.shutdownNow();
			// Preserve interrupt status
			Thread.currentThread().interrupt();
		}
	}
}
